{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pickle\n",
    "import csv\n",
    "import os\n",
    "import torch\n",
    "from torch_geometric.data import Data\n",
    "from tqdm import tqdm\n",
    "np.random.seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../input/yoochoose-clicks.p','rb') as f:\n",
    "    df = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>session_id</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>item_id</th>\n",
       "      <th>category</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>2014-04-07T10:51:09.277Z</td>\n",
       "      <td>214536502</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>2014-04-07T10:54:09.868Z</td>\n",
       "      <td>214536500</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>2014-04-07T10:54:46.998Z</td>\n",
       "      <td>214536506</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>2014-04-07T10:57:00.306Z</td>\n",
       "      <td>214577561</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2</td>\n",
       "      <td>2014-04-07T13:56:37.614Z</td>\n",
       "      <td>214662742</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   session_id                 timestamp    item_id category\n",
       "0           1  2014-04-07T10:51:09.277Z  214536502        0\n",
       "1           1  2014-04-07T10:54:09.868Z  214536500        0\n",
       "2           1  2014-04-07T10:54:46.998Z  214536506        0\n",
       "3           1  2014-04-07T10:57:00.306Z  214577561        0\n",
       "4           2  2014-04-07T13:56:37.614Z  214662742        0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>session_id</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>item_id</th>\n",
       "      <th>price</th>\n",
       "      <th>quantity</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>420374</td>\n",
       "      <td>2014-04-06T18:44:58.314Z</td>\n",
       "      <td>214537888</td>\n",
       "      <td>12462</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>420374</td>\n",
       "      <td>2014-04-06T18:44:58.325Z</td>\n",
       "      <td>214537850</td>\n",
       "      <td>10471</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>281626</td>\n",
       "      <td>2014-04-06T09:40:13.032Z</td>\n",
       "      <td>214535653</td>\n",
       "      <td>1883</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>420368</td>\n",
       "      <td>2014-04-04T06:13:28.848Z</td>\n",
       "      <td>214530572</td>\n",
       "      <td>6073</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>420368</td>\n",
       "      <td>2014-04-04T06:13:28.858Z</td>\n",
       "      <td>214835025</td>\n",
       "      <td>2617</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   session_id                 timestamp    item_id  price  quantity\n",
       "0      420374  2014-04-06T18:44:58.314Z  214537888  12462         1\n",
       "1      420374  2014-04-06T18:44:58.325Z  214537850  10471         1\n",
       "2      281626  2014-04-06T09:40:13.032Z  214535653   1883         1\n",
       "3      420368  2014-04-04T06:13:28.848Z  214530572   6073         1\n",
       "4      420368  2014-04-04T06:13:28.858Z  214835025   2617         1"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "buy_df = pd.read_csv('../input/yoochoose-buys.dat', header=None)\n",
    "buy_df.columns=['session_id','timestamp','item_id','price','quantity']\n",
    "buy_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "session_id     4431931\n",
       "timestamp     24590089\n",
       "item_id          48255\n",
       "category           331\n",
       "dtype: int64"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# filter out item session with length < 2\n",
    "df['valid_session'] = df.session_id.map(df.groupby('session_id')['item_id'].size() > 2)\n",
    "df = df.loc[df.valid_session].drop('valid_session',axis=1)\n",
    "df.nunique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "session_id    100000\n",
       "timestamp     555860\n",
       "item_id        22875\n",
       "category         136\n",
       "dtype: int64"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# #randomly sample a couple of them\n",
    "sampled_session_id = np.random.choice(df.session_id.unique(), 100000, replace=False)\n",
    "df = df.loc[df.session_id.isin(sampled_session_id)]\n",
    "df.nunique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5.5589"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# average length of session \n",
    "df.groupby('session_id')['item_id'].size().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>session_id</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>item_id</th>\n",
       "      <th>category</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>341</th>\n",
       "      <td>108</td>\n",
       "      <td>2014-04-03T11:29:34.491Z</td>\n",
       "      <td>5446</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>342</th>\n",
       "      <td>108</td>\n",
       "      <td>2014-04-03T11:30:31.064Z</td>\n",
       "      <td>5447</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>343</th>\n",
       "      <td>108</td>\n",
       "      <td>2014-04-03T11:31:36.145Z</td>\n",
       "      <td>5447</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>616</th>\n",
       "      <td>194</td>\n",
       "      <td>2014-04-04T14:29:16.071Z</td>\n",
       "      <td>15871</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>617</th>\n",
       "      <td>194</td>\n",
       "      <td>2014-04-04T14:29:47.845Z</td>\n",
       "      <td>15871</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     session_id                 timestamp  item_id  category\n",
       "341         108  2014-04-03T11:29:34.491Z     5446         0\n",
       "342         108  2014-04-03T11:30:31.064Z     5447         0\n",
       "343         108  2014-04-03T11:31:36.145Z     5447         0\n",
       "616         194  2014-04-04T14:29:16.071Z    15871         0\n",
       "617         194  2014-04-04T14:29:47.845Z    15871         0"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.preprocessing import LabelEncoder\n",
    "\n",
    "item_encoder = LabelEncoder()\n",
    "category_encoder = LabelEncoder()\n",
    "df['item_id'] = item_encoder.fit_transform(df.item_id )\n",
    "df['category']= category_encoder.fit_transform(df.category.apply(str))\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>session_id</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>item_id</th>\n",
       "      <th>price</th>\n",
       "      <th>quantity</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>101</th>\n",
       "      <td>420508</td>\n",
       "      <td>2014-04-07T15:37:26.815Z</td>\n",
       "      <td>10020</td>\n",
       "      <td>1046</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>102</th>\n",
       "      <td>420508</td>\n",
       "      <td>2014-04-07T15:37:26.896Z</td>\n",
       "      <td>7907</td>\n",
       "      <td>1570</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>103</th>\n",
       "      <td>420508</td>\n",
       "      <td>2014-04-07T15:37:26.950Z</td>\n",
       "      <td>15606</td>\n",
       "      <td>627</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>106</th>\n",
       "      <td>281747</td>\n",
       "      <td>2014-04-06T10:32:23.948Z</td>\n",
       "      <td>16075</td>\n",
       "      <td>523</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>107</th>\n",
       "      <td>281747</td>\n",
       "      <td>2014-04-06T10:32:23.984Z</td>\n",
       "      <td>16075</td>\n",
       "      <td>523</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     session_id                 timestamp  item_id  price  quantity\n",
       "101      420508  2014-04-07T15:37:26.815Z    10020   1046         1\n",
       "102      420508  2014-04-07T15:37:26.896Z     7907   1570         1\n",
       "103      420508  2014-04-07T15:37:26.950Z    15606    627         1\n",
       "106      281747  2014-04-06T10:32:23.948Z    16075    523         2\n",
       "107      281747  2014-04-06T10:32:23.984Z    16075    523         1"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "buy_df = buy_df.loc[buy_df.session_id.isin(df.session_id)]\n",
    "buy_df['item_id'] = item_encoder.transform(buy_df.item_id)\n",
    "buy_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{664: [7370],\n",
       " 3074: [16928, 13598, 11040, 11040],\n",
       " 3852: [10119],\n",
       " 4436: [1057],\n",
       " 4759: [15490, 15717],\n",
       " 4982: [9725, 10495],\n",
       " 6223: [4810],\n",
       " 6839: [4665],\n",
       " 10836: [15658, 2136],\n",
       " 12593: [5822],\n",
       " 12874: [13835, 13826],\n",
       " 13159: [4806],\n",
       " 13332: [15530],\n",
       " 13726: [9367],\n",
       " 14314: [13862, 15596, 15526, 15714],\n",
       " 14819: [15719, 15493],\n",
       " 18166: [7034],\n",
       " 18577: [7486],\n",
       " 18744: [50],\n",
       " 21441: [269],\n",
       " 22427: [16103, 16088],\n",
       " 23912: [13831],\n",
       " 25232: [10435],\n",
       " 26399: [16117, 16155, 16112, 16145, 16277, 16143, 17974],\n",
       " 27157: [12373, 17125],\n",
       " 27814: [8985],\n",
       " 28444: [4814],\n",
       " 29269: [15666],\n",
       " 30136: [12175, 12925, 12904, 14183],\n",
       " 30416: [12672],\n",
       " 31394: [13831],\n",
       " 34594: [15824, 3479, 12876, 12874, 12876, 14185],\n",
       " 34787: [13862, 13827],\n",
       " 35436: [15720, 15438],\n",
       " 35903: [15588, 14698],\n",
       " 38431: [15823, 15872],\n",
       " 39811: [15823, 15827],\n",
       " 41699: [13837],\n",
       " 42681: [18395, 16075],\n",
       " 44522: [5446],\n",
       " 45246: [15669, 15659],\n",
       " 45301: [16091, 16110],\n",
       " 46324: [5573],\n",
       " 46489: [1397],\n",
       " 47369: [9352],\n",
       " 51357: [16079],\n",
       " 51544: [15823, 15825, 16075],\n",
       " 51772: [16075, 16052, 16026],\n",
       " 52164: [15596, 15524],\n",
       " 53476: [167],\n",
       " 53976: [15827, 15823],\n",
       " 54169: [15513, 15720],\n",
       " 54807: [15669],\n",
       " 56414: [13891],\n",
       " 57128: [16075, 15823],\n",
       " 59013: [14690],\n",
       " 59756: [8331],\n",
       " 60524: [18264, 17976, 18424, 17974, 3480, 11666],\n",
       " 60999: [16076, 15825],\n",
       " 61868: [2408, 6300],\n",
       " 63334: [14498, 14498],\n",
       " 64253: [15535, 16075],\n",
       " 65071: [15669, 14135, 14185],\n",
       " 69871: [15605, 15606, 14399],\n",
       " 71478: [16362, 16030, 16031],\n",
       " 72772: [15826],\n",
       " 73206: [4814, 4834],\n",
       " 73473: [16896],\n",
       " 73986: [12885],\n",
       " 75451: [13829, 8017, 13601],\n",
       " 75814: [14046],\n",
       " 76606: [10188, 16146, 16145, 17143],\n",
       " 79177: [4814, 5822, 16103, 12901],\n",
       " 79298: [1326],\n",
       " 79326: [10248, 16362],\n",
       " 80404: [350],\n",
       " 81699: [16356, 15543, 15810],\n",
       " 82056: [10119, 957],\n",
       " 85031: [12482],\n",
       " 85543: [12381],\n",
       " 85778: [11334],\n",
       " 86762: [13193],\n",
       " 88529: [15829, 15829],\n",
       " 89621: [13826, 13826],\n",
       " 89752: [15826, 15823],\n",
       " 90551: [15826],\n",
       " 91802: [2485],\n",
       " 93167: [18395],\n",
       " 93288: [3930, 8549, 2920, 413],\n",
       " 93767: [14305, 19156, 19407, 14009],\n",
       " 98044: [15658, 15659, 15658, 15659],\n",
       " 101083: [18395],\n",
       " 101799: [15539, 16356, 18527, 27, 12177, 7003],\n",
       " 102328: [18395, 16035],\n",
       " 103563: [13862],\n",
       " 103808: [16035, 16129],\n",
       " 104433: [18504],\n",
       " 105521: [13849, 13828],\n",
       " 105869: [13848, 13848],\n",
       " 109477: [16025, 7350],\n",
       " 110102: [4998, 4998],\n",
       " 110794: [12375],\n",
       " 111327: [12751],\n",
       " 112347: [14213, 15513, 15378],\n",
       " 112654: [16438],\n",
       " 117724: [16928, 16929],\n",
       " 117768: [9680],\n",
       " 120052: [16050, 16129, 16056, 15824],\n",
       " 120882: [15826, 15823, 15825],\n",
       " 120913: [16035, 15825, 16076],\n",
       " 121068: [13847],\n",
       " 121143: [10170],\n",
       " 121344: [1214],\n",
       " 122206: [15827, 15826],\n",
       " 123721: [15662, 15661],\n",
       " 124086: [6386],\n",
       " 124227: [10119, 10119],\n",
       " 125127: [15823, 15825],\n",
       " 125616: [15669, 15659],\n",
       " 126351: [11136],\n",
       " 127773: [18600, 18395, 18395, 18600],\n",
       " 129689: [5842],\n",
       " 129791: [4950],\n",
       " 130951: [16079],\n",
       " 134273: [10119],\n",
       " 137139: [16079],\n",
       " 137747: [18474],\n",
       " 138412: [10163, 5822],\n",
       " 138417: [15823, 16082, 16025, 16026, 8637],\n",
       " 141537: [16129, 16035, 16129],\n",
       " 142303: [15792],\n",
       " 144766: [16110, 16034, 16110, 16034],\n",
       " 144921: [14031],\n",
       " 145081: [955, 4894],\n",
       " 146792: [9066],\n",
       " 149698: [13852, 15548, 3927, 3931],\n",
       " 150246: [12906, 13973],\n",
       " 150567: [12724],\n",
       " 150571: [1128, 15714, 16052],\n",
       " 151086: [5874, 5874],\n",
       " 152851: [17234, 17234],\n",
       " 153399: [6538, 6590],\n",
       " 154192: [13847, 15525],\n",
       " 154284: [15446, 15539],\n",
       " 154452: [15645, 12672],\n",
       " 154518: [2567],\n",
       " 154589: [1397],\n",
       " 156792: [15824, 12481],\n",
       " 158172: [448, 4006, 4338, 18527, 2136, 4338, 448, 18527, 2136, 4006],\n",
       " 159307: [18680, 16026],\n",
       " 162448: [15783, 15667],\n",
       " 164473: [13848],\n",
       " 165179: [18892],\n",
       " 165883: [18395],\n",
       " 167734: [16026,\n",
       "  16102,\n",
       "  16103,\n",
       "  16075,\n",
       "  16129,\n",
       "  16035,\n",
       "  16035,\n",
       "  16103,\n",
       "  16075,\n",
       "  16026,\n",
       "  16102,\n",
       "  16129],\n",
       " 168091: [5804, 6799],\n",
       " 170533: [16106, 15733],\n",
       " 171901: [2557, 15660],\n",
       " 172539: [15823, 16075],\n",
       " 176701: [7262],\n",
       " 178473: [15587, 18394, 17338, 16089, 14129, 12153, 695, 17125],\n",
       " 179402: [15601, 3733],\n",
       " 180324: [17],\n",
       " 181049: [16147, 16147],\n",
       " 181247: [1519],\n",
       " 181903: [13827, 13827, 13863],\n",
       " 181956: [12198],\n",
       " 182532: [489],\n",
       " 183897: [4784],\n",
       " 184067: [8087],\n",
       " 184964: [13851],\n",
       " 185072: [15658, 15669],\n",
       " 186181: [16052, 5123, 12961],\n",
       " 188948: [3263, 2533, 4017],\n",
       " 190192: [6224, 914, 12715, 5123, 13767, 12732, 16096],\n",
       " 191042: [16050, 19531],\n",
       " 191361: [14699],\n",
       " 191846: [16362],\n",
       " 192402: [5743],\n",
       " 193379: [15823, 15826],\n",
       " 193656: [15783, 5127],\n",
       " 193919: [16152, 15830],\n",
       " 194862: [13847],\n",
       " 195702: [17594],\n",
       " 195879: [18657, 14601],\n",
       " 197229: [2495],\n",
       " 198592: [16069],\n",
       " 198623: [10386, 16165],\n",
       " 200388: [10352],\n",
       " 202996: [13862, 13862],\n",
       " 203307: [13829],\n",
       " 204636: [12672],\n",
       " 204824: [12715],\n",
       " 205628: [13863, 13863],\n",
       " 206533: [11353],\n",
       " 207317: [5446],\n",
       " 207474: [16896, 14698],\n",
       " 207726: [15530],\n",
       " 210548: [16769, 13829],\n",
       " 210747: [16052, 16026, 15733, 16075],\n",
       " 211159: [215],\n",
       " 211247: [13829, 13852],\n",
       " 211604: [10219],\n",
       " 211956: [16033, 17974, 18264, 16033, 17974, 18264, 16033, 17974, 18264],\n",
       " 212643: [15590],\n",
       " 212659: [12277, 12719],\n",
       " 216117: [16056, 16090],\n",
       " 217563: [16026, 16052],\n",
       " 221456: [30],\n",
       " 224541: [5446],\n",
       " 226372: [16026, 16052],\n",
       " 229713: [13828],\n",
       " 229921: [2568, 2561, 3177, 10093],\n",
       " 231216: [15712],\n",
       " 231408: [12925],\n",
       " 232867: [14185],\n",
       " 233008: [16440],\n",
       " 234518: [13852],\n",
       " 234644: [14739, 14738, 14687],\n",
       " 234892: [13829, 13829],\n",
       " 236396: [18620],\n",
       " 236594: [14129, 19504, 2849, 15823, 15826],\n",
       " 237323: [2244, 12174],\n",
       " 238218: [15812, 15810, 15548],\n",
       " 238548: [14129, 14678, 16097, 7162],\n",
       " 238876: [15823, 16075, 16056, 15792],\n",
       " 242079: [12490],\n",
       " 242136: [13050],\n",
       " 242316: [14601, 8753, 16159, 18395],\n",
       " 243112: [18657],\n",
       " 244526: [15823, 16075, 14213, 5422, 13522],\n",
       " 246082: [4114, 4114],\n",
       " 249331: [7208],\n",
       " 251018: [16082, 714, 16129, 16025, 16051, 15826, 16025, 16074],\n",
       " 251511: [15858],\n",
       " 251596: [15528],\n",
       " 251766: [16090, 10190, 16025, 8605, 15869, 3139, 8607],\n",
       " 252796: [13193, 13192],\n",
       " 253521: [2561],\n",
       " 254476: [15824, 16056, 16035],\n",
       " 256631: [46],\n",
       " 256888: [16026],\n",
       " 258976: [11754],\n",
       " 259308: [3382],\n",
       " 259367: [16840, 16843, 16840, 16843],\n",
       " 259448: [16034],\n",
       " 261459: [13848, 13848],\n",
       " 261488: [10163],\n",
       " 262583: [18395, 18395],\n",
       " 262886: [16117, 12904, 15490, 8872, 3265, 15715, 8874, 13992, 16866, 16874],\n",
       " 262982: [18759],\n",
       " 263577: [6757, 13853, 13853],\n",
       " 263593: [15823, 16076, 15827, 15825],\n",
       " 264424: [3759, 11966],\n",
       " 265111: [14682],\n",
       " 269976: [13862],\n",
       " 270578: [15649],\n",
       " 270901: [12381, 8364, 19232, 10619],\n",
       " 275372: [15443],\n",
       " 276693: [15669],\n",
       " 277677: [14681, 14739],\n",
       " 277917: [13863],\n",
       " 278768: [10193, 17143, 16117, 16143, 16143],\n",
       " 278893: [1475, 13932],\n",
       " 279613: [14185],\n",
       " 279876: [15490, 14554, 15606],\n",
       " 281747: [16075, 16075, 3919, 16026, 13784],\n",
       " 281889: [481, 6974, 483],\n",
       " 282629: [17132],\n",
       " 283526: [9679],\n",
       " 285879: [16106, 16106],\n",
       " 285926: [13828],\n",
       " 286098: [13485, 8565, 13486, 13488, 9500, 9428, 13487],\n",
       " 287644: [15723, 4814, 15826, 15723],\n",
       " 288194: [15821, 16152, 15821, 16152, 15821, 16152, 15821, 16152],\n",
       " 288469: [13826],\n",
       " 288598: [2811],\n",
       " 288758: [14124, 9121, 14078, 14124, 14078, 9121],\n",
       " 289939: [810],\n",
       " 290443: [15823, 18395],\n",
       " 290924: [15867],\n",
       " 292056: [16951],\n",
       " 293234: [15724, 2561],\n",
       " 294064: [16896, 3958],\n",
       " 296098: [19313,\n",
       "  10661,\n",
       "  10728,\n",
       "  10703,\n",
       "  10703,\n",
       "  10661,\n",
       "  10728,\n",
       "  19313,\n",
       "  10703,\n",
       "  19313,\n",
       "  10728,\n",
       "  10661,\n",
       "  10661,\n",
       "  19313,\n",
       "  10728,\n",
       "  10661,\n",
       "  10728,\n",
       "  10703,\n",
       "  10703,\n",
       "  19313],\n",
       " 296514: [2988],\n",
       " 299038: [10248, 16362],\n",
       " 299093: [12704],\n",
       " 299311: [5446],\n",
       " 301073: [4538, 15783],\n",
       " 301144: [11305, 15860, 15810],\n",
       " 304679: [13147],\n",
       " 304773: [15606, 15587],\n",
       " 305097: [13862, 13864, 13862],\n",
       " 307207: [13852],\n",
       " 307402: [13853, 13835, 13836],\n",
       " 310213: [13852, 16929],\n",
       " 311778: [13825],\n",
       " 312012: [16356, 15714],\n",
       " 313419: [1214, 4189],\n",
       " 313631: [16157],\n",
       " 314311: [11948, 2054, 13862, 13862],\n",
       " 315249: [16896],\n",
       " 318108: [10475],\n",
       " 319651: [16079],\n",
       " 321926: [5446, 15669],\n",
       " 322699: [13862, 13826],\n",
       " 326788: [13827, 13862],\n",
       " 328261: [16896],\n",
       " 329126: [9736],\n",
       " 329332: [16129, 16035, 8753],\n",
       " 332248: [12373, 7336, 18997, 19050],\n",
       " 333853: [13845],\n",
       " 334398: [18759],\n",
       " 335847: [4780, 4780],\n",
       " 337112: [16051, 16106, 16129],\n",
       " 337392: [12435, 16152, 15821, 12138],\n",
       " 338086: [17234, 17234],\n",
       " 338623: [13124, 13124],\n",
       " 339276: [13848],\n",
       " 339839: [15541, 13145],\n",
       " 340817: [10020],\n",
       " 341411: [15715, 15587, 4321],\n",
       " 342279: [10169, 12878, 16090],\n",
       " 343774: [13862],\n",
       " 344909: [15720],\n",
       " 344959: [5842],\n",
       " 345243: [16896],\n",
       " 346413: [15848, 15834],\n",
       " 346617: [13850],\n",
       " 347289: [15807],\n",
       " 347538: [19227],\n",
       " 348349: [8656],\n",
       " 349938: [9765, 13081],\n",
       " 350484: [6816],\n",
       " 352222: [14115, 10535, 2963],\n",
       " 352786: [5894, 5666],\n",
       " 352956: [12381],\n",
       " 353063: [15659, 15669],\n",
       " 353346: [14739, 14687, 15826, 16051, 16026],\n",
       " 358676: [16090, 16035, 16035, 16090],\n",
       " 359259: [15649, 15645, 12672],\n",
       " 359267: [13886, 4632, 14698],\n",
       " 359888: [18759, 16893],\n",
       " 360363: [14127],\n",
       " 360582: [15490, 14554, 15717, 15831, 15715],\n",
       " 362744: [18703, 4795, 16106],\n",
       " 363044: [18657],\n",
       " 363842: [14213, 15378],\n",
       " 364578: [15826, 8753],\n",
       " 365153: [18857],\n",
       " 365414: [13828, 13828],\n",
       " 365967: [15446, 2518],\n",
       " 367438: [15823, 15826],\n",
       " 367902: [15029],\n",
       " 368144: [5842, 5843],\n",
       " 368747: [13862],\n",
       " 371494: [16074],\n",
       " 372653: [12370, 12381, 15714],\n",
       " 372736: [11281],\n",
       " 372949: [15530],\n",
       " 376931: [10768, 10766, 10743, 12707, 10741],\n",
       " 377099: [13863],\n",
       " 377233: [4129],\n",
       " 379359: [1397, 1397],\n",
       " 379927: [11956],\n",
       " 382842: [13830, 13830],\n",
       " 383092: [8651],\n",
       " 383533: [15821, 16152, 15821],\n",
       " 383806: [5180],\n",
       " 385497: [8408],\n",
       " 387293: [15782, 15661, 14555],\n",
       " 388241: [12414],\n",
       " 389174: [18620],\n",
       " 389609: [5075],\n",
       " 392554: [5097],\n",
       " 397726: [15490, 14213, 15719, 16095, 16146, 14514, 15720],\n",
       " 397814: [13778],\n",
       " 398391: [12961, 3766, 17171],\n",
       " 398791: [15827, 16051, 16026],\n",
       " 400239: [11491],\n",
       " 401476: [14683],\n",
       " 402341: [15669, 15658, 15659],\n",
       " 402427: [2562],\n",
       " 402999: [13849],\n",
       " 404724: [15511],\n",
       " 406521: [13829, 13829],\n",
       " 408091: [16362, 10248],\n",
       " 408213: [873],\n",
       " 408509: [12181],\n",
       " 410546: [13831, 13831, 13831],\n",
       " 412157: [13831],\n",
       " 412679: [16331, 16088, 16103],\n",
       " 413379: [16818, 16813, 16824, 16814, 16820],\n",
       " 413856: [1397],\n",
       " 415947: [14219, 4539, 18395],\n",
       " 416779: [18657],\n",
       " 417343: [16927],\n",
       " 417844: [7516],\n",
       " 418542: [15869, 12373, 13713, 16438, 13710],\n",
       " 418774: [12370, 15814],\n",
       " 420508: [10020, 7907, 15606],\n",
       " 420612: [2568, 3301],\n",
       " 421681: [14308, 4463],\n",
       " 421939: [14698, 15588, 14184],\n",
       " 422691: [15716, 13680, 15609],\n",
       " 423222: [15723, 15723],\n",
       " 424134: [15658, 15658],\n",
       " 425214: [18657],\n",
       " 427489: [18395, 8753, 8753, 18395, 18395, 8753, 18395, 8753],\n",
       " 427792: [5122],\n",
       " 430931: [16082, 15823],\n",
       " 433862: [2560],\n",
       " 434294: [5084],\n",
       " 435912: [14684, 14740],\n",
       " 437794: [12381, 15815],\n",
       " 439189: [13854],\n",
       " 439823: [13073],\n",
       " 439856: [7459],\n",
       " 440388: [16050,\n",
       "  16075,\n",
       "  16927,\n",
       "  19504,\n",
       "  16033,\n",
       "  16025,\n",
       "  3479,\n",
       "  18370,\n",
       "  16129,\n",
       "  16056,\n",
       "  16090],\n",
       " 440741: [720, 18789, 2329],\n",
       " 443116: [18815, 18854, 15606, 15782, 6096],\n",
       " 443273: [18811, 19060, 10721],\n",
       " 444321: [16896, 15548, 15810],\n",
       " 446062: [10968],\n",
       " 446611: [10728],\n",
       " 447311: [11281, 15715],\n",
       " 447974: [18630, 18631],\n",
       " 448201: [18395],\n",
       " 452081: [16082, 16103, 16129],\n",
       " 452433: [3805],\n",
       " 454514: [13826],\n",
       " 454988: [15662, 13714],\n",
       " 455769: [13975, 12750, 12897, 15814, 8923],\n",
       " 455897: [17508],\n",
       " 456894: [14610],\n",
       " 458597: [15525],\n",
       " 458707: [215],\n",
       " 459436: [18263, 18038, 13490],\n",
       " 459796: [19288,\n",
       "  18789,\n",
       "  10660,\n",
       "  18792,\n",
       "  18786,\n",
       "  18781,\n",
       "  18789,\n",
       "  18786,\n",
       "  19288,\n",
       "  10660,\n",
       "  18792,\n",
       "  18781],\n",
       " 460161: [16129],\n",
       " 460502: [1735],\n",
       " 460531: [12791],\n",
       " 461829: [10402],\n",
       " 468754: [15830, 15792],\n",
       " 468904: [10124, 4134],\n",
       " 469487: [14699, 14740],\n",
       " 472021: [18395, 18600, 18602],\n",
       " 472233: [16025, 15825, 16075],\n",
       " 473271: [3862],\n",
       " 476339: [16120, 16029, 16119, 16118, 15831, 16142],\n",
       " 476729: [15669, 15659],\n",
       " 477824: [11258],\n",
       " 478877: [16075, 13862],\n",
       " 478947: [5051],\n",
       " 479894: [16145, 16117, 17143, 16143, 7347],\n",
       " 480899: [15446, 14061],\n",
       " 481006: [15669, 15659, 5849, 18113],\n",
       " 481982: [14109, 16033],\n",
       " 483777: [10020],\n",
       " 483836: [16103, 16034, 16092, 16034, 16092, 16103],\n",
       " 484351: [15658],\n",
       " 484992: [3808],\n",
       " 485566: [15705, 18631],\n",
       " 485833: [16929],\n",
       " 485942: [483, 18527],\n",
       " 486211: [14683, 14739],\n",
       " 487977: [14683, 14739, 14699],\n",
       " 489383: [16277],\n",
       " 489664: [18395],\n",
       " 492316: [7469, 4898, 4895],\n",
       " 492871: [1397],\n",
       " 494056: [18759],\n",
       " 494507: [15814, 15806, 15806, 15814],\n",
       " 495018: [15793],\n",
       " 496381: [4439, 16345],\n",
       " 499207: [17739],\n",
       " 500657: [15834, 4573],\n",
       " 500846: [16110, 15873],\n",
       " 504338: [16148, 4968],\n",
       " 504357: [10020, 14740],\n",
       " 505033: [7321],\n",
       " 505206: [4784, 14700],\n",
       " 506106: [18621, 18060, 15712],\n",
       " 506751: [15832, 4557, 4637, 15835, 15833],\n",
       " 507074: [7629],\n",
       " 507803: [6375, 2500],\n",
       " 507808: [15823],\n",
       " 509067: [19216, 19377, 19329],\n",
       " 509744: [18264],\n",
       " 511633: [14554, 14399, 10534, 15606],\n",
       " 511896: [16147, 10260, 12140],\n",
       " 512031: [4495, 13547, 16156, 4495, 13547, 16156],\n",
       " 512392: [17976, 17973, 17974, 18264],\n",
       " 512547: [15823, 15825],\n",
       " 514349: [5842],\n",
       " 515577: [15724, 17974, 12874, 12876, 16034, 10173, 17966],\n",
       " 516284: [12350, 15251],\n",
       " 517847: [6379],\n",
       " 518358: [2101, 2101],\n",
       " 522229: [19370],\n",
       " 522423: [14609, 14677],\n",
       " 523074: [17973, 11666, 15716, 14555],\n",
       " 523246: [17966, 17974],\n",
       " 523764: [18395],\n",
       " 523947: [13848, 13848, 17456],\n",
       " 524331: [14685],\n",
       " 525718: [6369],\n",
       " 526524: [16110],\n",
       " 526541: [14179, 14177, 4314, 5225, 12199],\n",
       " 527432: [14127, 13052],\n",
       " 529928: [18657],\n",
       " 532223: [16068, 5447],\n",
       " 532227: [17975],\n",
       " 532301: [11340, 9413, 16440],\n",
       " 532708: [17131],\n",
       " 533373: [5874],\n",
       " 534032: [15649, 16839],\n",
       " 534418: [41],\n",
       " 534628: [17783, 17815],\n",
       " 534937: [11205, 16313, 15489, 15808],\n",
       " 535658: [15832, 15833],\n",
       " 535997: [16029, 16030],\n",
       " 536277: [15823, 15827],\n",
       " 536482: [16934],\n",
       " 536579: [4644, 4644],\n",
       " 537258: [13053],\n",
       " 540604: [17975],\n",
       " 540669: [4784],\n",
       " 541392: [11105],\n",
       " 545451: [4784],\n",
       " 545711: [17389, 1774],\n",
       " 546756: [14699, 14699],\n",
       " 552893: [14687, 14739],\n",
       " 553333: [16030, 16029, 16032],\n",
       " 554193: [363],\n",
       " 556344: [10248, 16362],\n",
       " 557162: [13547, 4499],\n",
       " 558272: [16440],\n",
       " 559582: [16437],\n",
       " 561018: [13791, 13790],\n",
       " 562718: [16928, 13852],\n",
       " 563354: [15724, 16026, 2976, 4931],\n",
       " 563716: [16147, 16149, 16281, 16363, 16967, 16364, 16156],\n",
       " 564222: [4495],\n",
       " 564543: [18421, 17974],\n",
       " 565626: [15832, 15848],\n",
       " 565692: [5787],\n",
       " 567318: [14134, 15792],\n",
       " 567601: [16148],\n",
       " 568542: [7459],\n",
       " 569314: [11337, 11337],\n",
       " 573521: [2181],\n",
       " 573856: [11118],\n",
       " 578242: [13073],\n",
       " 578999: [16155, 13052, 13053, 18601],\n",
       " 579296: [6268],\n",
       " 579943: [15814],\n",
       " 580002: [16052],\n",
       " 583411: [13862],\n",
       " 585047: [5842],\n",
       " 585319: [2972, 16363],\n",
       " 585432: [16112, 15720],\n",
       " 585671: [15800, 13944, 13504, 15724],\n",
       " 587392: [10085, 2561, 2545],\n",
       " 587494: [4338],\n",
       " 587973: [3251],\n",
       " 588077: [15832, 15838, 16091, 15833, 16110],\n",
       " 589599: [17337],\n",
       " 589803: [16149],\n",
       " 590574: [19812],\n",
       " 590733: [3919, 16274, 16274, 2792],\n",
       " 591052: [15670, 15670],\n",
       " 596624: [15849],\n",
       " 597506: [14699, 12874, 12878, 12900],\n",
       " 598349: [16277, 16112, 3366, 13792, 3888, 3888, 13783, 13777, 13801],\n",
       " 599966: [7508, 16082],\n",
       " 600164: [15661, 12672],\n",
       " 601273: [12901, 12901],\n",
       " 601291: [17462, 17401],\n",
       " 601296: [15781, 15781, 15834, 15724, 13992],\n",
       " 606364: [16912],\n",
       " 607118: [15826, 16082],\n",
       " 607241: [15583],\n",
       " 608912: [15832],\n",
       " 610456: [16148],\n",
       " 611276: [4911],\n",
       " 612124: [10119],\n",
       " 614178: [5874],\n",
       " 614619: [8692, 8692],\n",
       " 614743: [12913, 12913],\n",
       " 615913: [16148, 12262, 15530],\n",
       " 616396: [15251],\n",
       " 619558: [4778],\n",
       " 620627: [6379, 4911],\n",
       " 622567: [15862, 15862],\n",
       " 622614: [17798],\n",
       " 626413: [16364, 12327, 16364, 16034],\n",
       " 627226: [18479],\n",
       " 627236: [15733, 15733, 15733, 15733],\n",
       " 628212: [16119],\n",
       " 630506: [10119],\n",
       " 631613: [15376],\n",
       " 631647: [16277, 16112],\n",
       " 631792: [16110, 16155, 15873, 16097],\n",
       " 634447: [6409],\n",
       " 636488: [17974, 17976, 16095, 15719, 16030],\n",
       " 636712: [16029],\n",
       " 637053: [12322],\n",
       " 639416: [12925, 12906, 14183, 15835],\n",
       " 639557: [15834, 16145, 15848],\n",
       " 639664: [15542],\n",
       " 641292: [16145],\n",
       " 644414: [2533, 4725],\n",
       " 646981: [15556, 2999],\n",
       " 647388: [17976, 17974, 18264],\n",
       " 647622: [15832, 15832],\n",
       " 647647: [1655, 12175],\n",
       " 651543: [15613],\n",
       " 652539: [16283, 17194, 16374],\n",
       " 652984: [5822],\n",
       " 654964: [7503, 16895, 14213],\n",
       " 656003: [12904],\n",
       " 656688: [18477],\n",
       " 660037: [96, 96],\n",
       " 660324: [8082],\n",
       " 660903: [15832, 15833, 18, 16031],\n",
       " 661101: [16230, 7164, 18657, 5333, 16155],\n",
       " 662401: [15378],\n",
       " 663599: [15834, 15833, 15835, 15832],\n",
       " 665514: [15823],\n",
       " 668204: [4439],\n",
       " 668378: [18759],\n",
       " 669973: [16361, 15721],\n",
       " 670422: [13827, 13827],\n",
       " 670452: [6962],\n",
       " 673892: [16345, 11435],\n",
       " 673996: [13147],\n",
       " 675592: [303],\n",
       " 676896: [15832],\n",
       " 677026: [14126, 15467, 15465],\n",
       " 678141: [5842],\n",
       " 678884: [6799, 6795, 6799, 6795],\n",
       " 679334: [16075, 16026],\n",
       " 680566: [10231],\n",
       " 681246: [15835, 15833],\n",
       " 681519: [15658, 7501],\n",
       " 685042: [13073],\n",
       " 685836: [16364, 16364, 16142],\n",
       " 685853: [1098],\n",
       " 689187: [15562, 15609],\n",
       " 689437: [9680, 6567],\n",
       " 691141: [13847, 13847, 13847],\n",
       " 691258: [15833],\n",
       " 693137: [15823, 15825],\n",
       " 694594: [17975,\n",
       "  13053,\n",
       "  14554,\n",
       "  16031,\n",
       "  3001,\n",
       "  16032,\n",
       "  16028,\n",
       "  13052,\n",
       "  16029,\n",
       "  16030,\n",
       "  16028,\n",
       "  16031,\n",
       "  14554,\n",
       "  17975,\n",
       "  16032,\n",
       "  13052,\n",
       "  16029,\n",
       "  16030,\n",
       "  3001,\n",
       "  13053],\n",
       " 695327: [13825],\n",
       " 698724: [15659],\n",
       " 698831: [16301, 1713, 19254, 18854],\n",
       " 699989: [14554, 15715, 14618],\n",
       " 701168: [4500, 16028, 18603, 16032, 18602],\n",
       " 702514: [150],\n",
       " 704844: [1397, 16281, 4979],\n",
       " 707079: [13502, 16144, 14518],\n",
       " 708996: [15837],\n",
       " 709946: [4521, 18603, 18601],\n",
       " 710129: [9679, 9679],\n",
       " 712441: [5822, 5822],\n",
       " 714477: [17976, 11666, 17976, 17973, 17974],\n",
       " 714578: [18550, 12311, 12314],\n",
       " 716798: [3139, 16277, 13836],\n",
       " 717916: [15823],\n",
       " 720164: [14061, 9773],\n",
       " 720237: [12935],\n",
       " 721713: [16829, 16828, 16834, 16837, 16835],\n",
       " 722209: [15827, 15823],\n",
       " 723513: [18264, 18421, 17966, 17973],\n",
       " 724086: [13830, 8379, 10130, 15717, 10082, 14633],\n",
       " 724593: [15438],\n",
       " 724928: [6409],\n",
       " 726321: [15867, 16149, 16149],\n",
       " 726348: [13832, 13833],\n",
       " 727131: [15719, 1525],\n",
       " 727516: [15779, 16109, 16108, 16108],\n",
       " 727804: [16155, 16277],\n",
       " 729984: [15811,\n",
       "  14109,\n",
       "  9314,\n",
       "  14105,\n",
       "  16364,\n",
       "  15378,\n",
       "  15811,\n",
       "  14109,\n",
       "  9314,\n",
       "  14105,\n",
       "  15378,\n",
       "  16364],\n",
       " 730713: [15832],\n",
       " 736773: [15667, 13748, 2895],\n",
       " 736911: [15835],\n",
       " 738569: [15823],\n",
       " 741408: [16110, 17966, 15873, 16097, 17975, 5039, 16091],\n",
       " 744246: [9066],\n",
       " 745382: [16052],\n",
       " 746962: [14683],\n",
       " 749009: [13992],\n",
       " 749659: [10707],\n",
       " 749741: [2219],\n",
       " 750643: [10017, 14215],\n",
       " 751232: [3888, 16112, 16277],\n",
       " 752201: [17976, 17975, 18264],\n",
       " 753936: [1230, 15808],\n",
       " 754146: [14183],\n",
       " 755537: [16146, 16030],\n",
       " 755763: [10169, 12925],\n",
       " 756253: [18395],\n",
       " 756632: [15446, 15539, 15438, 1094, 3476],\n",
       " 757204: [400],\n",
       " 758347: [16147],\n",
       " 758356: [15832, 15833],\n",
       " 760552: [13791, 13790, 16112, 3139],\n",
       " 761127: [16277, 16112],\n",
       " 761247: [14008],\n",
       " 761901: [9721, 14555, 16897],\n",
       " 762883: [14124],\n",
       " 763356: [10294, 10248],\n",
       " 764059: [13829],\n",
       " 765322: [1091],\n",
       " 766978: [4784, 14700],\n",
       " 767024: [16152, 16152],\n",
       " 767657: [4036],\n",
       " 768677: [9680, 13975],\n",
       " 769198: [6002, 17986, 17985, 8771, 4923, 14489, 16, 6001],\n",
       " 769557: [18704],\n",
       " 770168: [1044],\n",
       " 771236: [11966],\n",
       " 771317: [9680, 9680],\n",
       " 772567: [18395, 18395],\n",
       " 775396: [13992, 14419, 14698],\n",
       " 776424: [19058, 18816],\n",
       " 778941: [15833, 15848, 16032, 16029],\n",
       " 779221: [10997, 13562],\n",
       " 780859: [15669],\n",
       " 782526: [15838, 15849],\n",
       " 782899: [17976, 17974, 18424, 18264],\n",
       " 783439: [15826, 15827],\n",
       " 785888: [15662, 17966, 17974, 1660],\n",
       " 786966: [10025, 2610],\n",
       " 787367: [9528],\n",
       " 787648: [1478, 15779, 12707],\n",
       " 788496: [16291, 16291],\n",
       " 788843: [15649],\n",
       " 789148: [16148, 16149],\n",
       " 793342: [171],\n",
       " 793603: [15251],\n",
       " 795971: [16147, 16029, 16155, 15848],\n",
       " 796943: [18759],\n",
       " 798881: [15659, 15669, 15657],\n",
       " 801239: [16110, 16029, 2784, 4333, 16110, 2784, 16029, 4333],\n",
       " 801338: [5163],\n",
       " 802001: [16277, 16155, 16031, 16277, 16031, 16155],\n",
       " 802257: [18781, 10671, 16147],\n",
       " 802992: [14685],\n",
       " 803757: [15873, 16095],\n",
       " 804597: [14060],\n",
       " 806356: [16147, 13791, 16145, 13790],\n",
       " 806542: [16894, 17974, 2894],\n",
       " 811819: [15733, 12350, 16030, 16029, 12404, 17976, 17966, 17974],\n",
       " 813497: [2915],\n",
       " 816324: [7459],\n",
       " 816928: [18525],\n",
       " 818982: [15658],\n",
       " 819381: [1397, 1397],\n",
       " 822608: [16967, 16282],\n",
       " 823092: [12901, 16149],\n",
       " 823891: [4367],\n",
       " 824838: [5054, 4358],\n",
       " 825356: [8395, 16028],\n",
       " 825851: [10248, 4499, 4499, 10248],\n",
       " 825968: [7549],\n",
       " 826473: [4544],\n",
       " 828227: [11948],\n",
       " 829834: [16313, 16147, 15829],\n",
       " 830721: [2502],\n",
       " 830923: [16, 816],\n",
       " 834477: [14179, 8887, 16035],\n",
       " 836676: [15545, 15588],\n",
       " 836782: [12672],\n",
       " 837333: [1675, 16362],\n",
       " 838163: [8786],\n",
       " 838919: [16318, 16032, 12901, 16029, 13052, 2675, 16364],\n",
       " 839164: [644, 13975],\n",
       " 839983: [3382, 3382],\n",
       " 840126: [16281, 16026, 16967, 16050, 168],\n",
       " 840518: [16155, 18],\n",
       " 841463: [5136],\n",
       " 841753: [5511, 5553],\n",
       " 842138: [12925, 12906, 14185],\n",
       " 844292: [4911, 15834, 4911, 15834],\n",
       " 844364: [15495, 15594],\n",
       " 844916: [9680, 15544],\n",
       " 847236: [16147],\n",
       " 847721: [4209],\n",
       " 850087: [15649],\n",
       " 850638: [14060],\n",
       " 851234: [10619, 18779, 18996],\n",
       " 851271: [996, 14051, 9308],\n",
       " 852043: [1043],\n",
       " 852113: [375],\n",
       " 852639: [18812, 18685, 518],\n",
       " 853907: [9679],\n",
       " 854249: [15835, 15832],\n",
       " 854514: [15823],\n",
       " 855487: [14687, 2495],\n",
       " 857101: [731, 15526, 11515, 14697, 663, 8002],\n",
       " 858303: [10168],\n",
       " 858989: [5842, 5843],\n",
       " 860744: [10631, 10660],\n",
       " 861072: [4784],\n",
       " 861572: [16026],\n",
       " 861682: [16026, 16052],\n",
       " 865203: [16362, 2054],\n",
       " 865864: [6409, 16149, 16147],\n",
       " 866957: [16799],\n",
       " 867847: [15832, 15833],\n",
       " 868947: [15530],\n",
       " 869576: [12350, 14774],\n",
       " 869609: [16155, 16117, 16143, 17143],\n",
       " 874004: [1397],\n",
       " 874896: [15823],\n",
       " 878082: [10620, 714],\n",
       " 878903: [17409],\n",
       " 880854: [16026, 16052],\n",
       " 881134: [19241, 18782],\n",
       " 882121: [16277, 16112],\n",
       " 882173: [15825, 15827],\n",
       " 882219: [18997, 19248, 19372],\n",
       " 882804: [17639],\n",
       " 883553: [15823, 15826],\n",
       " 883941: [10653, 19229, 19231, 10680],\n",
       " 886262: [16117, 16145],\n",
       " 887327: [15800, 15800, 15800],\n",
       " 887523: [9189, 15029, 15230, 9185],\n",
       " 888758: [12922, 14185],\n",
       " 888983: [7517],\n",
       " 889477: [16147, 16064, 16064, 16147],\n",
       " 892556: [16052, 4784],\n",
       " 893112: [5781, 2245],\n",
       " 894392: [9680, 17130],\n",
       " 896981: [16149, 5278],\n",
       " 898271: [16311],\n",
       " 899893: [4118, 14605, 16032],\n",
       " 900303: [17974, 16147],\n",
       " 900999: [3923],\n",
       " 901178: [400],\n",
       " 901638: [12825, 12825],\n",
       " 902306: [7668],\n",
       " 902709: [1044],\n",
       " 903146: [16088, 16095, 16147, 16097, 14697],\n",
       " 903229: [15823, 15823],\n",
       " 903844: [15823],\n",
       " 903929: [16030],\n",
       " 903971: [15835],\n",
       " 904331: [15838],\n",
       " 905107: [11687, 11687],\n",
       " 905133: [2664, 19056, 19354],\n",
       " 911523: [15838],\n",
       " 912124: [11256],\n",
       " 916676: [9679, 9679],\n",
       " 916996: [16096, 14182],\n",
       " 919731: [17537],\n",
       " 920076: [12935],\n",
       " 920668: [4310, 4310],\n",
       " 921007: [13862],\n",
       " 923772: [18936, 18936],\n",
       " 923798: [7232, 18657],\n",
       " 924419: [16438, 15733, 16052],\n",
       " 924459: [5874],\n",
       " 925007: [16103, 16088, 16104, 16102],\n",
       " 925631: [14140],\n",
       " 926282: [16149, 16148, 15867],\n",
       " 928473: [19479],\n",
       " 929968: [11364],\n",
       " 930078: [13621],\n",
       " 931178: [6816],\n",
       " 931813: [16152, 15251, 16157, 16362],\n",
       " 932911: [14683],\n",
       " 933367: [13852],\n",
       " 934587: [13832],\n",
       " 935052: [15717, 15715],\n",
       " 935256: [4340, 4340],\n",
       " 936817: [13655],\n",
       " 937507: [16361],\n",
       " 937534: [13928, 13927],\n",
       " 937789: [16364, 18456],\n",
       " 938139: [4838, 3555],\n",
       " 938354: [4454],\n",
       " 939312: [14739, 15446],\n",
       " 941438: [15849, 19555, 6369, 15833],\n",
       " 941462: [15779],\n",
       " 945021: [5110],\n",
       " 945877: [15838],\n",
       " 946856: [18395],\n",
       " 947676: [4074],\n",
       " 950218: [2796],\n",
       " 950774: [16364],\n",
       " 951412: [14409, 13862],\n",
       " 952632: [58],\n",
       " 956253: [16363, 19277, 10605, 19509, 10926],\n",
       " 956526: [16443, 16281, 16442],\n",
       " 956576: [14044, 3664],\n",
       " 956894: [19232, 19057, 19058],\n",
       " 958026: [1204],\n",
       " 958464: [14609],\n",
       " 961132: [200, 16365, 17456],\n",
       " 961306: [16118, 16118],\n",
       " 961767: [8314, 4340, 15834, 7954, 15848, 7410],\n",
       " 963128: [16302, 16370, 17429, 17569],\n",
       " 963254: [19555, 6380],\n",
       " 963398: [5446],\n",
       " 963943: [5175],\n",
       " 965807: [13547, 15781, 4495, 4499, 13547, 4499, 4495, 15781],\n",
       " 966657: [9679],\n",
       " 966999: [18395, 18600, 18603, 18602],\n",
       " 967788: [13854, 13854],\n",
       " 968437: [4544],\n",
       " 970472: [16896, 16896],\n",
       " 971377: [5175],\n",
       " 971676: [13848],\n",
       " 974012: [13547, 4495],\n",
       " 975234: [15848, 10445],\n",
       " 981486: [15329, 16274, 15778, 15870],\n",
       " 984226: [16364, 18456, 15780, 15779, 16373],\n",
       " 985657: [7554, 7554],\n",
       " 986043: [18148],\n",
       " 987003: [13192, 13193, 18527],\n",
       " 987906: [19490, 12216],\n",
       " 988716: [2054, 2300],\n",
       " 989837: [2406],\n",
       " 991161: [10178, 16352, 6083],\n",
       " 992053: [16311, 14677, 19232],\n",
       " 995948: [16148, 16149],\n",
       " 997079: [11353],\n",
       " 999487: [16967, 16299, 13958, 16369],\n",
       " 1001549: [14418, 13053],\n",
       " 1003829: [16364],\n",
       " 1004047: [4793, 4795],\n",
       " 1005814: [18997, 19397],\n",
       " 1008569: [16148, 12824, 16147, 12824, 16148, 16147],\n",
       " 1009607: [15832, 15848],\n",
       " 1011039: [15835, 15848],\n",
       " 1012652: [5896, 1396],\n",
       " 1013098: [18395],\n",
       " 1014288: [16376, 5742, 17149, 17124],\n",
       " 1015566: [14928, 19611, 14928, 19611, 14928],\n",
       " 1015642: [4648, 4521, 4645],\n",
       " 1015966: [16312],\n",
       " 1016117: [4199],\n",
       " 1018137: [12300, 11592],\n",
       " 1019316: [2103],\n",
       " 1024571: [200, 10613],\n",
       " 1025914: [15832, 19561],\n",
       " 1026343: [12196, 12134, 16096],\n",
       " 1027459: [16443, 16318, 16442, 16442, 16281],\n",
       " 1027499: [15832],\n",
       " 1029092: [7500],\n",
       " 1029377: [10117, 15587, 12874, 16279],\n",
       " 1032497: [19880],\n",
       " 1032537: [11283],\n",
       " 1034551: [16149, 16148],\n",
       " 1035063: [18009, 8361, 16142],\n",
       " 1035522: [16318, 15846],\n",
       " 1038488: [4916, 16355],\n",
       " 1039187: [19563, 6379, 2502],\n",
       " 1039621: [3113],\n",
       " 1039871: [15826, 15823],\n",
       " 1040589: [7266, 16031],\n",
       " 1040711: [18395],\n",
       " 1040791: [19462],\n",
       " 1040904: [13862, 15715],\n",
       " 1043561: [7190],\n",
       " 1044062: [18778],\n",
       " 1044853: [16442, 16281],\n",
       " 1044972: [2466],\n",
       " 1047248: [6015, 9448, 12177],\n",
       " 1047341: [11055],\n",
       " 1047444: [16365],\n",
       " 1051241: [2554, 2552],\n",
       " 1051937: [14078],\n",
       " 1052293: [16283],\n",
       " 1056519: [16914],\n",
       " 1057601: [17234],\n",
       " 1058669: [16076, 5830],\n",
       " 1059554: [17966],\n",
       " 1061114: [18563, 15443],\n",
       " 1062209: [12134, 11967, 12134, 11967],\n",
       " 1062383: [17542],\n",
       " 1064421: [216],\n",
       " 1067004: [5923, 17149, 15839],\n",
       " 1067233: [15810, 15549],\n",
       " 1067744: [5039, 3265],\n",
       " 1069413: [16147, 15823],\n",
       " 1071853: [3701],\n",
       " 1072362: [16318, 16318, 17149, 15715],\n",
       " 1073517: [6901, 7846],\n",
       " 1074236: [10119],\n",
       " 1077563: [14681, 14683, 14739],\n",
       " 1078531: [15837],\n",
       " 1078674: [15849, 16148, 16148, 15849],\n",
       " 1080936: [16147],\n",
       " 1082764: [17066, 17065],\n",
       " ...}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "buy_item_dict = dict(buy_df.groupby('session_id')['item_id'].apply(list))\n",
    "buy_item_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>session_id</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>item_id</th>\n",
       "      <th>category</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>341</th>\n",
       "      <td>108</td>\n",
       "      <td>2014-04-03T11:29:34.491Z</td>\n",
       "      <td>5446</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>342</th>\n",
       "      <td>108</td>\n",
       "      <td>2014-04-03T11:30:31.064Z</td>\n",
       "      <td>5447</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>343</th>\n",
       "      <td>108</td>\n",
       "      <td>2014-04-03T11:31:36.145Z</td>\n",
       "      <td>5447</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>616</th>\n",
       "      <td>194</td>\n",
       "      <td>2014-04-04T14:29:16.071Z</td>\n",
       "      <td>15871</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>617</th>\n",
       "      <td>194</td>\n",
       "      <td>2014-04-04T14:29:47.845Z</td>\n",
       "      <td>15871</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     session_id                 timestamp  item_id  category\n",
       "341         108  2014-04-03T11:29:34.491Z     5446         0\n",
       "342         108  2014-04-03T11:30:31.064Z     5447         0\n",
       "343         108  2014-04-03T11:31:36.145Z     5447         0\n",
       "616         194  2014-04-04T14:29:16.071Z    15871         0\n",
       "617         194  2014-04-04T14:29:47.845Z    15871         0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch_geometric.data import InMemoryDataset\n",
    "from tqdm import tqdm\n",
    "\n",
    "class YooChooseDataset(InMemoryDataset):\n",
    "    def __init__(self, root, transform=None, pre_transform=None):\n",
    "        super(YooChooseDataset, self).__init__(root, transform, pre_transform)\n",
    "        self.data, self.slices = torch.load(self.processed_paths[0])\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        return []\n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        return ['../input/yoochoose_click_binary_100000_sess.dataset']\n",
    "\n",
    "    def download(self):\n",
    "        pass\n",
    "    \n",
    "    def process(self):\n",
    "        \n",
    "        data_list = []\n",
    "\n",
    "        # process by session_id\n",
    "        grouped = df.groupby('session_id')\n",
    "        for session_id, group in tqdm(grouped):\n",
    "            le = LabelEncoder()\n",
    "            sess_item_id = le.fit_transform(group.item_id)\n",
    "            group = group.reset_index(drop=True)\n",
    "            group['sess_item_id'] = sess_item_id\n",
    "            node_features = group.loc[group.session_id==session_id,['sess_item_id','item_id','category']].sort_values('sess_item_id')[['item_id','category']].drop_duplicates().values\n",
    "\n",
    "            node_features = torch.LongTensor(node_features).unsqueeze(1)\n",
    "            target_nodes = group.sess_item_id.values[1:]\n",
    "            source_nodes = group.sess_item_id.values[:-1]\n",
    "\n",
    "            edge_index = torch.tensor([source_nodes,\n",
    "                                   target_nodes], dtype=torch.long)\n",
    "            x = node_features\n",
    "\n",
    "            if session_id in buy_item_dict:\n",
    "                positive_indices = le.transform(buy_item_dict[session_id])\n",
    "                label = np.zeros(len(node_features))\n",
    "                label[positive_indices] = 1\n",
    "            else:\n",
    "                label = [0] * len(node_features)\n",
    "\n",
    "\n",
    "            y = torch.FloatTensor(label)\n",
    "\n",
    "            data = Data(x=x, edge_index=edge_index, y=y)\n",
    "\n",
    "            data_list.append(data)\n",
    "        \n",
    "        data, slices = self.collate(data_list)\n",
    "        torch.save((data, slices), self.processed_paths[0])\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataset = YooChooseDataset('../')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# grouped = df.groupby('session_id')\n",
    "# data_list= []\n",
    "# for session_id, group in tqdm(grouped):\n",
    "#     le = LabelEncoder()\n",
    "#     sess_item_id = le.fit_transform(group.item_id)\n",
    "#     group = group.reset_index(drop=True)\n",
    "#     group['sess_item_id'] = sess_item_id\n",
    "#     node_features = group.loc[group.session_id==session_id,['sess_item_id','item_id']].sort_values('sess_item_id').item_id.drop_duplicates().values\n",
    "\n",
    "#     node_features = torch.LongTensor(node_features).unsqueeze(1)\n",
    "#     target_nodes = group.sess_item_id.values[1:]\n",
    "#     source_nodes = group.sess_item_id.values[:-1]\n",
    "\n",
    "#     edge_index = torch.tensor([source_nodes,\n",
    "#                            target_nodes], dtype=torch.long)\n",
    "#     x = node_features\n",
    "    \n",
    "#     if session_id in buy_item_dict:\n",
    "#         positive_indices = le.transform(buy_item_dict[session_id])\n",
    "#         label = np.zeros(len(node_features))\n",
    "#         label[positive_indices] = 1\n",
    "#     else:\n",
    "#         label = [0] * len(node_features)\n",
    "\n",
    "    \n",
    "#     y = torch.FloatTensor(label)\n",
    "    \n",
    "#     data = Data(x=x, edge_index=edge_index, y=y)\n",
    "    \n",
    "#     data_list.append(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(80000, 10000, 10000)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = dataset.shuffle()\n",
    "one_tenth_length = int(len(dataset) * 0.1)\n",
    "train_dataset = dataset[:one_tenth_length * 8]\n",
    "val_dataset = dataset[one_tenth_length*8:one_tenth_length * 9]\n",
    "test_dataset = dataset[one_tenth_length*9:]\n",
    "len(train_dataset), len(val_dataset), len(test_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.data import DataLoader\n",
    "batch_size= 512\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size)\n",
    "val_loader = DataLoader(val_dataset, batch_size=batch_size)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(22875, 135)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_items = df.item_id.max() +1\n",
    "num_categories = df.category.max()+1\n",
    "num_items , num_categories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_dim = 128\n",
    "from torch_geometric.nn import GraphConv, TopKPooling, GatedGraphConv, SAGEConv, SGConv\n",
    "from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp\n",
    "import torch.nn.functional as F\n",
    "class Net(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        self.conv1 = GraphConv(embed_dim * 2, 128)\n",
    "        self.pool1 = TopKPooling(128, ratio=0.9)\n",
    "        self.conv2 = GraphConv(128, 128)\n",
    "        self.pool2 = TopKPooling(128, ratio=0.9)\n",
    "        self.conv3 = GraphConv(128, 128)\n",
    "        self.pool3 = TopKPooling(128, ratio=0.9)\n",
    "        self.item_embedding = torch.nn.Embedding(num_embeddings=num_items, embedding_dim=embed_dim)\n",
    "        self.category_embedding = torch.nn.Embedding(num_embeddings=num_categories, embedding_dim=embed_dim)        \n",
    "        self.lin1 = torch.nn.Linear(256, 256)\n",
    "        self.lin2 = torch.nn.Linear(256, 128)\n",
    "        self.bn1 = torch.nn.BatchNorm1d(128)\n",
    "        self.bn2 = torch.nn.BatchNorm1d(64)\n",
    "        self.act1 = torch.nn.ReLU()\n",
    "        self.act2 = torch.nn.ReLU()        \n",
    "  \n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "        \n",
    "        item_id = x[:,:,0]\n",
    "        category = x[:,:,1]\n",
    "        \n",
    "\n",
    "        emb_item = self.item_embedding(item_id).squeeze(1)\n",
    "        emb_category = self.category_embedding(category).squeeze(1)\n",
    "        \n",
    "#         emb_item = emb_item.squeeze(1)\n",
    "#         emb_cat\n",
    "        x = torch.cat([emb_item, emb_category], dim=1)  \n",
    "#         print(x.shape)\n",
    "        x = F.relu(self.conv1(x, edge_index))\n",
    "#                 print(x.shape)\n",
    "        x, edge_index, _, batch, _ = self.pool1(x, edge_index, None, batch)\n",
    "        x1 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)\n",
    "\n",
    "        x = F.relu(self.conv2(x, edge_index))\n",
    "     \n",
    "        x, edge_index, _, batch, _ = self.pool2(x, edge_index, None, batch)\n",
    "        x2 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)\n",
    "\n",
    "        x = F.relu(self.conv3(x, edge_index))\n",
    "\n",
    "        x, edge_index, _, batch, _ = self.pool3(x, edge_index, None, batch)\n",
    "        x3 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)\n",
    "\n",
    "        x = x1 + x2 + x3\n",
    "\n",
    "        x = self.lin1(x)\n",
    "        x = self.act1(x)\n",
    "        x = self.lin2(x)\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.act2(x)      \n",
    "        \n",
    "        outputs = []\n",
    "        for i in range(x.size(0)):\n",
    "            output = torch.matmul(emb_item[data.batch == i], x[i,:])\n",
    "\n",
    "            outputs.append(output)\n",
    "              \n",
    "        x = torch.cat(outputs, dim=0)\n",
    "        x = torch.sigmoid(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')\n",
    "model = Net().to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "crit = torch.nn.BCELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    model.train()\n",
    "\n",
    "    loss_all = 0\n",
    "    for data in train_loader:\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "\n",
    "        label = data.y.to(device)\n",
    "        loss = crit(output, label)\n",
    "        loss.backward()\n",
    "        loss_all += data.num_graphs * loss.item()\n",
    "        optimizer.step()\n",
    "    return loss_all / len(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import roc_auc_score\n",
    "def evaluate(loader):\n",
    "    model.eval()\n",
    "\n",
    "    predictions = []\n",
    "    labels = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for data in loader:\n",
    "\n",
    "            data = data.to(device)\n",
    "            pred = model(data).detach().cpu().numpy()\n",
    "\n",
    "            label = data.y.detach().cpu().numpy()\n",
    "            predictions.append(pred)\n",
    "            labels.append(label)\n",
    "\n",
    "    predictions = np.hstack(predictions)\n",
    "    labels = np.hstack(labels)\n",
    "    \n",
    "    return roc_auc_score(labels, predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001, Loss: 0.70601, Train Auc: 0.51948, Val Auc: 0.52077, Test Auc: 0.51932\n",
      "Epoch: 002, Loss: 0.67905, Train Auc: 0.54509, Val Auc: 0.55940, Test Auc: 0.53070\n",
      "Epoch: 003, Loss: 0.61984, Train Auc: 0.55491, Val Auc: 0.54600, Test Auc: 0.53157\n",
      "Epoch: 004, Loss: 0.58205, Train Auc: 0.56236, Val Auc: 0.55520, Test Auc: 0.53699\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-161-f38f140dc697>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m200\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m     \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m     \u001b[0mtrain_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mevaluate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_loader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m     \u001b[0mval_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mevaluate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mval_loader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0mtest_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mevaluate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtest_loader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-160-fd7ffda79879>\u001b[0m in \u001b[0;36mevaluate\u001b[0;34m(loader)\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 12\u001b[0;31m             \u001b[0mpred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     13\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m             \u001b[0mlabel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    487\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    488\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 489\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    490\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    491\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-157-e16bbc0e6def>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, data)\u001b[0m\n\u001b[1;32m     38\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrelu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_index\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     39\u001b[0m \u001b[0;31m#                 print(x.shape)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 40\u001b[0;31m         \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpool1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     41\u001b[0m         \u001b[0mx1\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mgmp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     42\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    487\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    488\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 489\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    490\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    491\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/py36/lib/python3.6/site-packages/torch_geometric/nn/pool/topk_pool.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, edge_index, edge_attr, batch)\u001b[0m\n\u001b[1;32m    104\u001b[0m         \u001b[0mscore\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    105\u001b[0m         \u001b[0mscore\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mscore\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mp\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 106\u001b[0;31m         \u001b[0mperm\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtopk\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mscore\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mratio\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    107\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mperm\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtanh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mscore\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mperm\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    108\u001b[0m         \u001b[0mbatch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mperm\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/py36/lib/python3.6/site-packages/torch_geometric/nn/pool/topk_pool.py\u001b[0m in \u001b[0;36mtopk\u001b[0;34m(x, ratio, batch)\u001b[0m\n\u001b[1;32m     29\u001b[0m     mask = [\n\u001b[1;32m     30\u001b[0m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlong\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mi\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 31\u001b[0;31m         \u001b[0mmax_num_nodes\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnum_nodes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     32\u001b[0m     ]\n\u001b[1;32m     33\u001b[0m     \u001b[0mmask\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmask\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/py36/lib/python3.6/site-packages/torch_geometric/nn/pool/topk_pool.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     29\u001b[0m     mask = [\n\u001b[1;32m     30\u001b[0m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlong\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mi\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 31\u001b[0;31m         \u001b[0mmax_num_nodes\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnum_nodes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     32\u001b[0m     ]\n\u001b[1;32m     33\u001b[0m     \u001b[0mmask\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmask\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "for epoch in range(1, 200):\n",
    "    loss = train()\n",
    "    train_acc = evaluate(train_loader)\n",
    "    val_acc = evaluate(val_loader)    \n",
    "    test_acc = evaluate(test_loader)\n",
    "    print('Epoch: {:03d}, Loss: {:.5f}, Train Auc: {:.5f}, Val Auc: {:.5f}, Test Auc: {:.5f}'.\n",
    "          format(epoch, loss, train_acc, val_acc, test_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
